package com.party.admin.web.controller.websocket;

import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;

/**
 * Created by admin on 2018/3/20.
 */

@Component
public class MyWebSocketHandler implements WebSocketHandler {

    Logger logger= LoggerFactory.getLogger(getClass());

    private static Map<String, Map<String, WebSocketSession>> sessionMap;

    private static Map<String, Map<String, String>> flagMap;

    static {
        sessionMap = Maps.newHashMap();
        flagMap = Maps.newHashMap();
    }

    /**
     * 建立连接后
     * @param webSocketSession
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
        String memberId = getMemberId(webSocketSession);
        String sessionId = getSessionId(webSocketSession);
        String fromFlag = getFromFlag(webSocketSession);
        if (null == memberId || null == sessionId) {
            return;
        }
        Map<String, WebSocketSession> map = sessionMap.get(memberId);
        Map<String, String> flag = flagMap.get(memberId);

        if (null == map) {
            Map<String, WebSocketSession> newMap = Maps.newHashMap();
            newMap.put(sessionId, webSocketSession);
            sessionMap.put(memberId, newMap);

            Map<String, String> newFlag = Maps.newHashMap();
            newFlag.put(sessionId, fromFlag);
            flagMap.put(memberId, newFlag);
            return;
        }

        WebSocketSession wss = map.get(sessionId);
        if (null == wss) {
            map.put(sessionId, webSocketSession);
            flag.put(sessionId, fromFlag);
            return;
        }

        map.remove(sessionId);
        map.put(sessionId, webSocketSession);

        flag.remove(sessionId);
        flag.put(sessionId, fromFlag);
    }

    /**
     * 消息处理
     * @param webSocketSession
     * @param webSocketMessage
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) throws Exception {

    }

    /**
     * 消息传输错误处理
     * @param webSocketSession
     * @param throwable
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) throws Exception {
        if (webSocketSession.isOpen()){
            webSocketSession.close();
        }
    }

    /**
     * 关闭连接后
     * @param webSocketSession
     * @param closeStatus
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) throws Exception {
        String memberId = null;
        String sessionId = null;
        try {
            memberId = getMemberId(webSocketSession);
            sessionId = getSessionId(webSocketSession);
            String fromFlag = getFromFlag(webSocketSession);
            if (null == memberId || null == sessionId) {
                logger.info("webSocket关闭异常, memberId = {}, sessionId = {}", memberId, sessionId);
                return;
            }
            if (null == sessionMap.get(memberId) || null == sessionMap.get(memberId).get(sessionId)) {
                logger.info("webSocket关闭异常, memberId = {}, sessionId = {}对应的用户没有找到webSocketSession",
                        memberId, sessionId);
                return;
            }
            Map<String, WebSocketSession> map = sessionMap.get(memberId);
            map.remove(sessionId);

            Map<String, String> flag = flagMap.get(memberId);
            flag.remove(sessionId);
            if (map.size() < 1) {
                sessionMap.remove(memberId);
                flagMap.remove(memberId);
            }
        } catch (Exception e) {
            logger.error("webSocket关闭异常, memberId = {}, sessionId = {}, 异常信息为{}", memberId, sessionId, e);
        }

    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 获取用户id
     * @param webSocketSession
     * @return
     */
    private String getMemberId(WebSocketSession webSocketSession) {
        try {
            return (String) webSocketSession.getAttributes().get(HandShake.CLIENT_ID);
        } catch (Exception e) {
            logger.error("webSocket链接时获取用户id失败", e);
            return null;
        }
    }

    /**
     * 获取用户session_id
     * @param webSocketSession
     * @return
     */
    private String getSessionId(WebSocketSession webSocketSession) {
        try {
            return (String) webSocketSession.getAttributes().get(HandShake.SESSION_ID);
        } catch (Exception e) {
            logger.error("webSocket链接时获取用户session_id失败", e);
            return null;
        }
    }

    private String getFromFlag(WebSocketSession webSocketSession) {
        try {
            return (String) webSocketSession.getAttributes().get(HandShake.FROM_FLAG);
        } catch (Exception e) {
            logger.error("webSocket链接时获取fromFlag失败", e);
            return null;
        }
    }

    /**
     * 消息推送
     * @param message 消息
     */
    public void send(final TextMessage message, String memberId, String sessionId){
        if (null == flagMap.get(memberId) || null == flagMap.get(memberId).get(sessionId) || !"true".equals(flagMap.get(memberId).get(sessionId))) {
            return;
        }
        try {
            if (null == sessionMap.get(memberId) || null == sessionMap.get(memberId).get(sessionId)) {
                logger.info("用户memberId = {}, sessionId = {}推送消息失败, webSocketSessionMap中没有找到该用户的session",
                        memberId, sessionId);
                return;
            }
            WebSocketSession session = sessionMap.get(memberId).get(sessionId);
            session.sendMessage(message);
        } catch (IOException e) {
            logger.info("发生消息异常", e);
        }
    }
}
